#
# messkit.tls_async_stream
#
# Copyright 2006 Helsinki Institute for Information Technology (HIIT)
# and the authors.
#
# Authors: Ken Rimey <rimey@hiit.fi>
#

# Permission is hereby granted, free of charge, to any person
# obtaining a copy of this software and associated documentation files
# (the "Software"), to deal in the Software without restriction,
# including without limitation the rights to use, copy, modify, merge,
# publish, distribute, sublicense, and/or sell copies of the Software,
# and to permit persons to whom the Software is furnished to do so,
# subject to the following conditions:
#
# The above copyright notice and this permission notice shall be
# included in all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
# EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
# MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
# IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY
# CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT,
# TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
# SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

"""
Stream implementation with support for TLS

This uses the tlslite (0.3.8) and cryptoID (0.3.3) libraries available
from http://trevp.net/ .
"""

from messkit.async_stream import Channel
from messkit.async_stream import Stream
from messkit.async_stream import StreamServer as _StreamServer
from messkit.async_stream import StreamManager as _StreamManager

__all__ = ['TLSStream', 'Stream', 'StreamServer', 'StreamManager']

def init_tls():
    try:
        TLSChannel
    except:
        _init_tls()

def _init_tls():
    global TLSChannel, TLSConnection, HandshakeSettings

    from tlslite.TLSConnection import TLSConnection
    from tlslite.HandshakeSettings import HandshakeSettings
    from messkit.TLSAsyncDispatcherMixIn import TLSAsyncDispatcherMixIn

    class HandshakeChannel(Channel):
        def __init__(self, map):
            Channel.__init__(self, map)
            self.handshake_callback = None

        def handle_connect(self):
            self.handle_handshake(error=False)

        def handle_error(self):
            self.handle_handshake(error=True)
            Channel.handle_error(self)

        def handle_handshake(self, error):
            cb = self.handshake_callback
            if cb is not None:
                self.handshake_callback = None
                cb(error)

    class TLSChannel(TLSAsyncDispatcherMixIn, HandshakeChannel):
        def __init__(self, map):
            HandshakeChannel.__init__(self, map)
            TLSAsyncDispatcherMixIn.__init__(self)

class TLSStream(Stream):
    def handshake(self, cb=None, role=None, settings=None, **params):
        assert self.channel.connected

        init_tls()

        if role is not None:
            assert role in ('client', 'server')
        elif self.initiating:
            role = 'client'
        else:
            role = 'server'
        self.tls_role = role

        if settings is None:
            settings = HandshakeSettings()
            settings.cipherNames.reverse() # Prefer fastest cipher (RC4).
        params['settings'] = settings

        def proxy(error):
            if not error:
                chain = self.get_cert_chain()
                self.public_key = chain.getEndEntityPublicKey()
                self.end_entity = self.public_key.hash()
            self.issue_callback(cb, error)

        self.schedule_write_handler(self.__handshake(proxy, params))

    def __handshake(self, cb, params):
        assert not self.rbuf.size, 'Input buffer not empty at start of TLS handshake.'
        assert not self.wbuf.size, 'Output buffer not empty at start of TLS handshake.'

        self.aggressive_write = False # Required for correctness.
        self.send_size = 3600 # Smaller value helps avoid write-write deadlocks.

        socket = self.channel.socket
        self.channel.stop()

        self.channel = TLSChannel(self.mgr.map)
        self.channel.connected = True
        self.channel.set_socket(socket)
        tls = self.channel.tlsConnection = TLSConnection(socket)

        if self.tls_role == 'client':
            handshaker = tls.handshakeClientCert(async=True, **params)
        else:
            handshaker = tls.handshakeServerAsync(reqCert=True, **params)

        self.channel.handshake_callback = cb
        self.channel.setHandshakeOp(handshaker)

        yield False

    def validate(self):
        self.cert_chain = self.get_cert_chain()
        if self.cert_chain.validate():
            self.cryptoID = self.cert_chain.cryptoID
            return True
        else:
            return False

    def get_cert_chain(self):
        session = self.channel.tlsConnection.session
        if self.tls_role == 'client':
            return session.serverCertChain
        else:
            return session.clientCertChain

class StreamServer(_StreamServer):
    pass

class StreamManager(_StreamManager):
    def _create_stream(self, *args):
        return TLSStream(self, *args)

    def _create_server(self, socket_params):
        return StreamServer(self, socket_params)
